
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import PointCloud2, PointField
import numpy as np
import time
import math
import config as cfg
from threading import Thread

class ROS2_Set():
    def __init__(self):
        rclpy.init(args=None)
        self.pointCloudPublisher = PointCloudPublisher()
    # 启动话题发布
    def launch_publisher(self):
        publisher_thread = Thread(target=self.__publisher_run)
        publisher_thread.daemon = True
        publisher_thread.start()
    
    def __publisher_run(self):
        while True:
            time.sleep(0.01)
            rclpy.spin_once(self.pointCloudPublisher)
        self.pointCloudPublisher.destroy_node()
        rclpy.shutdown()
        
    def set_data(self, laser_sensor_data):
        self.pointCloudPublisher.set_data(laser_sensor_data)
        
    
class PointCloudPublisher(Node):
    def __init__(self):
        super().__init__('standalone_pcl_publisher')
        self.publisher_ = self.create_publisher(PointCloud2, cfg.LASER_POINT_ROS2_TOPOIC_NAME, 10)
        timer_period = 1.0 / cfg.LASER_POINT_ROS2_PUBLISH_FREQ
        self.timer = self.create_timer(timer_period, self.timer_callback)
        
        # 生成初始示例数据
        self.num_points = 90
        self.yaw = np.random.uniform(-np.pi, np.pi, self.num_points)    # 偏航角
        self.pitch = np.random.uniform(-np.pi/4, np.pi/4, self.num_points) # 俯仰角
        self.distance = np.random.uniform(1, 10, self.num_points)       # 距离

    def set_data(self, laser_sensor_data):
        self.num_points = len(laser_sensor_data["range"])
        self.yaw = np.array([laser_sensor_data["pos"] % (2 * math.pi)] * self.num_points)
        self.pitch = np.array([x * 0.03491 for x in list(range(0, 90))])
        self.distance = np.array(laser_sensor_data["range"])
        # 处理无效点
        invalid_mask = (self.distance == -1)
        self.distance[invalid_mask] = np.nan  # 标记无效点

    def timer_callback(self):
        def height_to_color(z, z_min, z_max):
            """
            将高度 z 映射到 RGB 颜色（0-255），颜色从蓝（低）渐变到红（高）
            """
            # 归一化高度到 [0, 1]
            if np.isnan(z) or np.isnan(z_min) or np.isnan(z_max):
                return 0   # 返回完全透明（Alpha=0，RGB=0
            else:
                normalized = (z - z_min) / (z_max - z_min)
                # 计算 RGB 分量
                r = int(255 * (1 - normalized))          # 红：从 0 到 255
                g = 0                              # 绿：固定为 0
                b = int(255 * normalized)    # 蓝：从 255 到 0
                # 返回 RGBA 格式（Alpha=255，表示不透明）
                return (r << 16) | (g << 8) | b | (255 << 24)  # 格式：0xAABBGGRR（RViz 默认解析方式）
            # 极坐标转笛卡尔坐标
        cos_pitch = np.cos(self.pitch)
        x = self.distance * cos_pitch * np.cos(self.yaw)
        y = self.distance * cos_pitch * np.sin(self.yaw)
        z = self.distance * np.sin(self.pitch)
            
        # 计算颜色（基于 Z 值
        z_no_nan = [x for x in z if not np.isnan(x)]
        z_min = np.min(z_no_nan)
        z_max = np.max(z_no_nan)
        colors = np.array([height_to_color(zi, z_min, z_max) for zi in z], dtype=np.uint32)
        # 合并坐标和颜色数据
        points = np.column_stack((x, y, z, colors)).astype(np.float32)  # 注意颜色字段需对齐偏移量

        # 构造PointCloud2消息
        msg = PointCloud2()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = cfg.POINTCLOUD2_FRAME_ID
        msg.height = 1
        msg.width = self.num_points
        msg.fields = [
            PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
            PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
            PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
            PointField(name='rgba', offset=12, datatype=PointField.UINT32, count=1),
        ]
        msg.is_bigendian = False
        msg.point_step = 16  # 3个float32 (12字节) + 1个uint32 (4字节) = 16字节
        msg.row_step = msg.point_step * msg.width
        msg.is_dense = True
        msg.data = points.tobytes()

        self.publisher_.publish(msg)